import numpy as np
import yaml
import os
import time
from algorithm.mathlib.utilities import PoseRobot
from pyquaternion import Quaternion
from algorithm.mathlib import utilities as utl


class Admittance:
    def __init__(self):
        main_path = os.path.dirname(os.path.dirname(os.path.dirname(os.path.realpath(__file__))))
        yaml_file = os.path.join(main_path, "config/AdmittanceParams.yaml")
        with open(yaml_file, encoding='utf-8') as f:
            content = f.read()
        data = yaml.load(content, Loader=yaml.FullLoader)
        # self._M = np.array(data["mass_arm"])
        # self._D = np.array(data["damping_arm"])
        # self._K = np.array(data["stiffness_arm"])
        self.error = np.zeros(6)
        self.twist_des = np.zeros(6)
        self.admittance_ratio = 1  # to adjust stiffness
        self.wrench_ext = np.zeros(6)
        self.label_time = time.time()

    def core(self, pose_des, pose_curr):
        if pose_des.qtn * pose_curr.qtn < 0.0:
            pose_curr.qtn = - pose_curr.qtn

        qtn_err = pose_curr.qtn * pose_des.qtn.inverse
        if qtn_err.norm > 0.001:
            qtn_err = qtn_err/qtn_err.norm
        self.error[3:] = qtn_err.axis * qtn_err.angle
        self.error[:3] = pose_curr.pos - pose_des.pos

        acc_des = np.linalg.inv(self._M) * (- self._K * self.error - self._D * self.twist_des + self.admittance_ratio * self.wrench_ext)

        # limit the acceleration
        acc_norm = np.linalg.norm(acc_des[:3])
        if acc_norm > self.acc_max:
            print("Admittance generates high arm acceleration!")
            acc_des[:3] *= self.acc_max / acc_norm

        self.label_time = time.time() - self.label_time

        self.twist_des += acc_des * self.label_time

        qtndot = 0.5 * utl.wedge4(self.twist_des[3:]) * pose_curr.qtn

        pose_des.pos = pose_curr.pos + self.twist_des[:3] * self.label_time
        pose_des.qtn = pose_curr.qtn + qtndot * self.label_time

        self.label_time = time.time()

        return self.pose_des




